#!/usr/bin/env python
import math
import sys
import  scipy.integrate as sciint
import numpy
from mpmath import *
import mpmath
import cmath

from unit import *
from bath import *
from constant import *
from propagator import *
from spectral_density import *
from tensor import *
from matrix_index_2d import *
from global_parameter import *




def generate_matrix_T_for_kmax1(T, J, stepsize, w0, env_parameters):
	M = [None for temp in range(0, 4)]
	for i in range(0, len(M)):
		M[i] = [None for temp in range(0, 4)]
	for i in range(0, len(M)):
		for j in range(0, len(M[0])):
			M[i][j] = propagator(generate_vector(i)[0], generate_vector(i)[1], generate_vector(j)[0], generate_vector(j)[1],w0, stepsize) * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, env_parameters, 0, generate_two_indice(j)[0], generate_two_indice(j)[1], generate_two_indice(j)[0], generate_two_indice(j)[1]) * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, env_parameters, 1, generate_two_indice(j)[0], generate_two_indice(j)[1], generate_two_indice(i)[0], generate_two_indice(i)[1])
	return matrix(M)

def generateT(tensor, s, i, N, w0, stepsize, T, J, env_parameters):
	for s[i] in range(0, 4):
		if i == (len(s) - 1):
			temp = 1
			for n in range(0, N):
				temp = temp * propagator(generate_vector(s[n + 1])[0], generate_vector(s[n + 1])[1], generate_vector(s[n])[0], generate_vector(s[n])[1], w0, stepsize) 
				for k in range(0, N + 1): 
					temp = temp * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, env_parameters, k, generate_two_indice(s[n])[0], generate_two_indice(s[n])[1], generate_two_indice(s[n + k])[0], generate_two_indice(s[n + k])[1])
			a = tensor
			for j in range( 0, len(s) - 1):
				a = a[s[j]]
			a[s[len(s) - 1]] = temp
		else:
			generateT(tensor, s, i + 1, N, w0, stepsize, T, Spectral_Density, env_parameters)


def generateT_Next(tensor, s, i, N, w0, stepsize, T, J, env_parameters):
	for s[i] in range(0, 4):
		if i == (len(s) - 1):
			temp = 1
			for n in range(0, N):
				temp = temp * propagator(generate_vector(s[n + 1])[0], generate_vector(s[n + 1])[1], generate_vector(s[n])[0], generate_vector(s[n])[1], w0, stepsize) 
				for k in range(0, N + 1):
					if n + k < 4 and n < 4:
						temp = temp * Feyman_Vernon_Specific_Distance_Contribution(T, J, stepsize, wc, kondon_coeff, k, generate_two_indice(s[n])[0], generate_two_indice(s[n])[1], generate_two_indice(s[n + k])[0], generate_two_indice(s[n + k])[1])
			a = tensor
			for j in range( 0, len(s) - 1):
				a = a[s[j]]
			a[s[len(s) - 1]] = temp
		else:
			generateT_Next(tensor, s, i + 1, N, w0, stepsize, T, Spectral_Density, env_parameters)




def cal_recursion(result, s, i, A, T): 
	for s[i] in range(0, 4):
		if i == 5:
			result[s[3]][s[4]][s[5]] += T[s[0]][s[1]][s[2]][s[3]][s[4]][s[5]] * A[s[0]][s[1]][s[2]]
		else:
			cal_recursion(result, s, i + 1, A, T)

def cal_recursion_for_result(result, s, i, A, T): 
	for s[i] in range(0, 4):
		if i == 3:
			result[s[3]] += T[s[0]][s[1]][s[2]][s[3]] * A[s[0]][s[1]][s[2]]
		else:
			cal_recursion_for_result(result, s, i + 1, A, T)
		

def calculate_for_result(T, A, N):
	result = generate_tensor(1, 0)
	s = [0 for temp in range(0, N + 1)]
	cal_recursion_for_result(result, s, 0, A, T)	 
	return result		

def calculate_for_kmax3(T, A, N):
	result = generate_tensor(N, 0)
	s = [0 for temp in range(0, N * 2)]
	cal_recursion(result, s, 0, A, T)	 
	return result		


def tensor_calculate_for_kmax1(w0, inital_state, temperature, Spectral_Density, step, env_parameters):
	A = matrix([0.5, 0.5, 0.5, 0.5])
	#A = matrix([1, 0, 0, 0])
	T = generate_matrix_T_for_kmax1(temperature, Spectral_Density, step,  w0, env_parameters)
	#print "0\t1\t0"
	print "0\t0.5\t0\t0.5"
	N = 1
	while N < 300: 
		A = T * A
		result = matrix([0, 0, 0, 0])	
		for i in range(0, len(result)):
			result[i] = A[i] * Feyman_Vernon_Specific_Distance_Contribution(temperature, Spectral_Density, step, env_parameters, 0, generate_two_indice(i)[0], generate_two_indice(i)[1], generate_two_indice(i)[0], generate_two_indice(i)[1])
		print "%lg\t%lf\t%lf\t%lf" % (1.0 * N * step, re(result[0]), re(result[0]) - re(result[1]), fabs(result[2]))
		N += 1

tensor_calculate_for_kmax1()
'''
def tensor_calculate_for_kmax(N = 3):
	sys.setrecursionlimit(10000)
	w0 = matrix([[0, 1], [1, 0]]) * delta0 + matrix([[1, 0], [0, -1]]) * epsi
	# for test
	step = 0.15 / delta0;
	#step = 0.000001e-9
	#generate tensor A
	A = generate_tensor_list(N, [1, 0, 0, 0])
	# generate null tensor T, just define the data structure
	T = generate_tensor(2 * N, None)
	T_Next = generate_tensor(N + 1, None)
	# generate the tensor simliar to vector
	Ts = [0 for i in range(0, 2 * N)]
	generateT(T, Ts, 0, N, w0, (1.0 / 3) * step, temperature, Spectral_Density, cutoff_freq, kondon)
	Ts = [0 for i in range(0, N + 1)]
	generateT_Next(T_Next, Ts, 0, N, w0, (1.0 / 3) * step, temperature, Spectral_Density, cutoff_freq, kondon)
	#generateT(T, Ts, 0, N, w0, 0.5e-15, 10 * hbar / k, Spectral_Density, cmtoInt(200000), 0.09)
	print "0\t1\t0"
	M = 1 
	while M < 40:
		result = calculate_for_result(T_Next, A, N)
		for i in range(0, len(result)):
			result[i] = result[i] * Feyman_Vernon_Specific_Distance_Contribution(temperature, Spectral_Density,  step, cutoff_freq, kondon, 0, generate_two_indice(i)[0], generate_two_indice(i)[1], generate_two_indice(i)[0], generate_two_indice(i)[1])
		print "%lg\t%lf\t%lf" % (1.0 * M * step, re(result[0]) - re(result[1]), fabs(result[2]))
		A = calculate_for_kmax3(T, A, N)
		M += 1
'''
